import sys
sys.path.append('/home/aistudio/ultralytics')
from ultralytics import YOLO

# 创建YOLO模型实例，加载预训练权重文件yolov8n.pt
model = YOLO('yolov8n.pt')

# //cpu   yolo train model=yolov8n.pt data=yolo-bvn.yaml epochs=30 workers=1 batch=16 device=cpu
# // gpu   yolo train model=yolov8n.pt data=yolo-bvn.yaml epochs=30 workers=1 batch=16 device=gpu
# Classes  yolo train model=yolov8n.pt data=yolo-bvn.yaml epochs=30 workers=1 batch=8 imgsz=416 device=gpu
# /home/aistudio/ultralytics/yolo-bvn.yaml
# 开始训练
model.train(
    data='/home/aistudio/ultralytics/yolo-bvn.yaml',  # 训练数据配置文件路径
    workers=0,  # 数据加载工作线程数
    epochs=500,  # 总训练轮数
    batch=16,  # 每次训练批次大小
    imgsz=640,  # 输入图像尺寸
    rect=False,  # 不进行矩形框校正
    resume=False,  # 不从之前的训练中断续训练
    device='cuda'  # 使用CUDA设备进行训练
)

